# 优化后的卷积神经网络
import torch
from torch import nn, optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
from PIL import Image
import matplotlib.pyplot as plt

# 一、导入手写数据集
train_dataset = datasets.MNIST(root='../data',
                               train=True,
                               transform=transforms.ToTensor(),
                               download=False)
test_dataset = datasets.MNIST(root='../data',
                              train=False,
                              transform=transforms.ToTensor(),
                              download=False)

# 二、定义数据包装器
batch_size = 64
train_loader = DataLoader(dataset=train_dataset,
                          batch_size=batch_size,
                          shuffle=True)
test_loader = DataLoader(dataset=test_dataset,
                         batch_size=batch_size,
                         shuffle=False)


# 三、定义模型:卷积网络层（无BatchNorm）
class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        self.conv1 = nn.Sequential(
            nn.Conv2d(1, 32, 3, 1, 1),  # 输入1通道，输出32通道，3x3卷积核
            nn.ReLU(),
            nn.MaxPool2d(2, 2)  # 28x28 -> 14x14
        )
        self.conv2 = nn.Sequential(
            nn.Conv2d(32, 64, 3, 1, 1),  # 32通道 -> 64通道
            nn.ReLU(),
            nn.MaxPool2d(2, 2)  # 14x14 -> 7x7
        )
        self.fc1 = nn.Sequential(
            nn.Linear(64 * 7 * 7, 256),
            nn.ReLU(),
            nn.Dropout(p=0.5)  # 保持dropout防止过拟合
        )
        self.fc2 = nn.Linear(256, 10)  # 输出10个类别

        # 初始化权重
        self._initialize_weights()

    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
        x = x.view(x.size(0), -1)  # 展平
        x = self.fc1(x)
        x = self.fc2(x)
        return x

    def _initialize_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.Linear):
                nn.init.normal_(m.weight, 0, 0.01)
                nn.init.constant_(m.bias, 0)


# 四、定义模型、损失函数和优化器
LR = 0.001
model = CNN()
loss_fn = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=LR)


# 五、训练函数
def model_train():
    model.train()
    total_loss = 0
    for inputs, targets in train_loader:
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = loss_fn(outputs, targets)
        loss.backward()
        optimizer.step()
        total_loss += loss.item()
    avg_loss = total_loss / len(train_loader)
    print(f"训练损失: {avg_loss:.4f}")
    return avg_loss


# 六、测试函数
def model_test():
    model.eval()
    total_correct = 0
    total_samples = 0
    with torch.no_grad():
        for inputs, targets in test_loader:
            outputs = model(inputs)
            _, predicted = torch.max(outputs, 1)
            total_correct += (predicted == targets).sum().item()
            total_samples += targets.size(0)
    accuracy = 100. * total_correct / total_samples
    print(f"测试集准确率: {accuracy:.2f}%")
    return accuracy


# 七、模型预测函数
def predict_image(model, img_path):
    # 图像预处理
    img = Image.open(img_path).resize((28, 28)).convert('L')
    img_tensor = transforms.ToTensor()(img)

    # 自动颜色校正（MNIST是黑底白字）
    if img_tensor.mean() > 0.5:
        img_tensor = 1 - img_tensor

    # 添加batch维度并预测
    img_tensor = img_tensor.unsqueeze(0)
    model.eval()
    with torch.no_grad():
        output = model(img_tensor)
        probs = torch.softmax(output, 1)[0]
        pred = torch.argmax(probs).item()

    # 输出结果
    print("\n预测结果详情:")
    print(f"预测数字: {pred}")
    print("各类别概率:")
    for i, p in enumerate(probs):
        print(f"数字{i}: {p.item():.2%}")

    # 可视化
    plt.imshow(img_tensor.squeeze(), cmap='gray')
    plt.title(f"Predicted: {pred}")
    plt.show()
    return pred


# 八、使用示例
if __name__ == "__main__":
    # 1. 训练模型
    # for epoch in range(10):
    #     print(f"\n第{epoch + 1}次训练开始...")
    #     model_train()
    #     model_test()

    # 2. 保存模型
    model_save_path = './model_save/mnist_cnn_net02.pth'
    # torch.save(model.state_dict(), model_save_path)
    # print(f"模型保存成功: {model_save_path}")

    # 3. 加载模型并预测

    model.load_state_dict(torch.load(model_save_path, weights_only=True))
    predict_image(model, "../images/7_1.png")
